#!/bin/env python
# -*- coding: utf-8 -*-
import taichi as ti
import numpy as np
import time
import sys
import os
from tqdm import tqdm

# 设置环境变量，确保GPU可以被充分利用
os.environ["TI_DEVICE_MEMORY_GB"] = "2.0"  # 分配2GB内存给GPU
os.environ["TI_ENABLE_CUDA"] = "1"         # 确保CUDA后端启用
os.environ["TI_USE_UNIFIED_MEMORY"] = "1"  # 使用统一内存

# 初始化Taichi，使用GPU后端
print("正在初始化Taichi (GPU)...")
ti.init(arch=ti.gpu, device_memory_GB=2.0, default_fp=ti.f32, default_ip=ti.i32,
        kernel_profiler=True, print_ir=False, print_kernel_llvm_ir=False)

# 打印Taichi配置信息
print(f"Taichi版本: {ti.__version__}")

# 创建大型数组以确保GPU有工作负载
n = 10000
print(f"创建 {n}x{n} 大小的数组...")

# 创建Taichi字段
x = ti.field(ti.f32, shape=(n, n))
y = ti.field(ti.f32, shape=(n, n))
z = ti.field(ti.f32, shape=(n, n))

# 定义一个用于测试GPU利用率的内核
@ti.kernel
def compute():
    for i, j in x:
        # 进行一些密集计算
        v = 0.0
        for k in range(50):  # 增加计算量
            v += ti.sin(ti.cast(i, ti.f32) * 0.1 + ti.cast(j, ti.f32) * 0.1) * ti.cos(ti.cast(k, ti.f32) * 0.1)
        z[i, j] = v

# 定义一个用于测试GPU内存带宽的内核
@ti.kernel
def memory_bandwidth():
    for i, j in x:
        y[i, j] = x[i, j] * 2.0

# 初始化数据
@ti.kernel
def init():
    for i, j in x:
        x[i, j] = ti.sin(ti.cast(i, ti.f32) * 0.1) * ti.cos(ti.cast(j, ti.f32) * 0.1)

# 持续运行以测试GPU利用率
def run_continuous_test(duration=10):
    print(f"持续测试GPU利用率 ({duration}秒)...")
    
    # 初始化数据
    init()
    
    # 预热
    print("预热中...")
    compute()
    ti.sync()
    
    # 连续运行计算内核
    start_time = time.time()
    iterations = 0
    
    with tqdm(total=duration, desc="GPU负载测试") as pbar:
        while time.time() - start_time < duration:
            compute()
            ti.sync()  # 确保GPU计算完成
            iterations += 1
            
            # 更新进度条
            elapsed = time.time() - start_time
            pbar.update(min(elapsed - pbar.n, duration - pbar.n))
    
    end_time = time.time()
    total_time = end_time - start_time
    ops_per_second = iterations / total_time
    
    print(f"完成 {iterations} 次迭代，每秒 {ops_per_second:.2f} 次操作")
    print(f"总耗时: {total_time:.4f} 秒")
    
    # 在这个时间段内，GPU应该显示较高的利用率
    print("测试期间，请检查您的GPU监控工具(如nvidia-smi、任务管理器等)观察GPU利用率")

# 测试内存带宽
def test_memory_bandwidth():
    print("测试GPU内存带宽...")
    
    # 初始化数据
    init()
    
    # 预热
    memory_bandwidth()
    ti.sync()
    
    # 测量内存带宽
    num_iterations = 100
    start_time = time.time()
    
    for _ in tqdm(range(num_iterations), desc="内存带宽测试"):
        memory_bandwidth()
        ti.sync()
    
    end_time = time.time()
    total_time = end_time - start_time
    
    # 计算带宽 (GB/s)
    # 每次操作读取x并写入y，共2个操作，每个元素4字节(float32)
    total_bytes = 2 * (n * n * 4) * num_iterations
    bandwidth_gb_per_sec = total_bytes / total_time / (1024 ** 3)
    
    print(f"内存带宽: {bandwidth_gb_per_sec:.2f} GB/s")

if __name__ == "__main__":
    # 测试GPU计算利用率
    run_continuous_test(duration=20)
    
    # 测试GPU内存带宽
    test_memory_bandwidth()
    
    # 输出性能分析信息
    if hasattr(ti, "profiler") and hasattr(ti.profiler, "print_kernel_profiler_info"):
        print("\n性能分析信息:")
        ti.profiler.print_kernel_profiler_info()
    elif hasattr(ti, "summary"):
        print("\n性能分析信息:")
        ti.summary()
    
    print("测试完成。GPU应该已经被充分利用。") 